{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\feature_extraction\\image.py:167: DeprecationWarning: `np.int` is a deprecated alias for the builtin `int`. To silence this warning, use `int` by itself. Doing this will not modify any behavior and is safe. When replacing `np.int`, you may wish to use e.g. `np.int64` or `np.int32` to specify the precision. If you wish to review your current use, check the release note link for additional information.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  dtype=np.int):\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:35: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:597: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps, copy_X=True, fit_path=True,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:836: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps, copy_X=True, fit_path=True,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:862: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps, positive=False):\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:1074: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  max_n_alphas=1000, n_jobs=1, eps=np.finfo(np.float).eps,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:1306: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  max_n_alphas=1000, n_jobs=1, eps=np.finfo(np.float).eps,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\least_angle.py:1442: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps, copy_X=True, positive=False):\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\randomized_l1.py:152: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  precompute=False, eps=np.finfo(np.float).eps,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\randomized_l1.py:318: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=np.finfo(np.float).eps, random_state=None,\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\sklearn\\linear_model\\randomized_l1.py:575: DeprecationWarning: `np.float` is a deprecated alias for the builtin `float`. To silence this warning, use `float` by itself. Doing this will not modify any behavior and is safe. If you specifically wanted the numpy scalar type, use `np.float64` here.\n",
      "Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations\n",
      "  eps=4 * np.finfo(np.float).eps, n_jobs=1,\n"
     ]
    }
   ],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.simplefilter('ignore')\n",
    "\n",
    "import os\n",
    "import cv2\n",
    "import importlib\n",
    "import torch\n",
    "import numpy as np\n",
    "import argparse\n",
    "import yaml\n",
    "from tqdm import tqdm\n",
    "from PIL import Image, ImageDraw, ImageOps\n",
    "import json\n",
    "import nltk\n",
    "import torchvision.transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F\n",
    "from scipy import signal\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.widgets import Button\n",
    "import clip\n",
    "import kornia\n",
    "import torchvision\n",
    "import skimage.color\n",
    "import random\n",
    "\n",
    "from utils_func import *\n",
    "from html_images import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--device', type=str, default='cuda:0')\n",
    "\n",
    "parser.add_argument('--root_dir', type=str, default='C:/MyFiles/CondTran/finals/quant')\n",
    "parser.add_argument('--log_dir', type=str, default='logs/bert_quant')\n",
    "parser.add_argument('--step', type=str, default='142124')\n",
    "\n",
    "# Load transformer\n",
    "args = parser.parse_args(args=[])\n",
    "os.chdir(args.root_dir)\n",
    "sys.path.append(args.root_dir)\n",
    "module = importlib.import_module(f'filltran.models.colorization')\n",
    "args.fill_model = getattr(module, 'Colorization')\n",
    "filltran = load_model(args.fill_model, args.log_dir, args.step).to(args.device).eval().requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\CondTran\\\\data\\\\raw\\\\in1.JPEG'\n",
    "I_color = Image.open(img_path).convert('RGB')\n",
    "I_gray = I_color.convert('L')\n",
    "idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "x_color = preprocess(I_color, None).to(args.device)\n",
    "x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "# Recon with original color indices\n",
    "'''idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "dec = output_to_pil(dec[0])\n",
    "dec.show()'''\n",
    "\n",
    "# Recon with full color indices\n",
    "x_color[:, 0, :, :] = 0.0; x_color[:, 1, :, :] = -0.6; x_color[:, 2, :, :] = -0.6\n",
    "idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "output_to_pil(dec[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\recon_r.png')\n",
    "output_to_pil(x_color[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\r.png')\n",
    "\n",
    "# Recon with full color indices\n",
    "x_color[:, 0, :, :] = -0.6; x_color[:, 1, :, :] = 0.0; x_color[:, 2, :, :] = -0.6\n",
    "idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "output_to_pil(dec[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\recon_g.png')\n",
    "output_to_pil(x_color[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\g.png')\n",
    "\n",
    "# Recon with full color indices\n",
    "x_color[:, 0, :, :] = -0.6; x_color[:, 1, :, :] = -0.6; x_color[:, 2, :, :] = 0.0\n",
    "idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "output_to_pil(dec[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\recon_b.png')\n",
    "output_to_pil(x_color[0]).save('C:\\\\MyFiles\\\\CondTran\\\\sample\\\\results\\\\b.png')\n",
    "\n",
    "# Recon with full gray gray indices\n",
    "'''I_allgray = 128 * np.ones(list(I_color.size[::-1])).astype(np.uint8)\n",
    "I_allgray = Image.fromarray(I_allgray)\n",
    "x_allgray = preprocess(I_allgray, None).to(args.device)\n",
    "idx_color, idx_allgray = filltran.hybrid_vqgan.encode(x_color, x_allgray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "idx_allgray = idx_allgray.reshape([idx_allgray.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, idx_allgray)\n",
    "dec = output_to_pil(dec[0])\n",
    "dec.show()'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gray reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "save_path = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\ablation\\\\quant_gray'\n",
    "img_size = [256, 256]\n",
    "\n",
    "files = os.listdir(img_path)\n",
    "\n",
    "for f in tqdm(files):\n",
    "    fname = f.split('.')[0]\n",
    "    I_color = Image.open(os.path.join(img_path, f)).convert('RGB').resize(img_size)\n",
    "    I_gray = I_color.convert('L')\n",
    "    idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "    x_color = preprocess(I_color, None).to(args.device)\n",
    "    x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "    # Recon with original color indices\n",
    "    idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "    idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "    idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "    dec_gt = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "    dec_gt = output_to_pil(dec_gt[0])\n",
    "\n",
    "    # Recon with full red color indices\n",
    "    I_rand = np.random.randint(0, 256, size=[1, 1, 3], dtype=np.uint8)\n",
    "    I_rand = I_rand.repeat(I_color.size[1], axis=0).repeat(I_color.size[0], axis=1)\n",
    "    I_rand = Image.fromarray(I_rand)\n",
    "    x_rand = preprocess(I_rand, None).to(args.device)\n",
    "    idx_rand, _ = filltran.hybrid_vqgan.encode(x_rand, None)\n",
    "    idx_rand = idx_rand.reshape([idx_rand.shape[0]]+idx_shape)\n",
    "    dec_rand = filltran.hybrid_vqgan.decode(idx_rand, idx_gray)\n",
    "    dec_rand = output_to_pil(dec_rand[0])\n",
    "\n",
    "    I_gray.convert('RGB').save( os.path.join(save_path, 'gray', f'{fname}.png') )\n",
    "    dec_gt.convert('L').convert('RGB').save( os.path.join(save_path, 'gt', f'{fname}.png') )\n",
    "    dec_rand.convert('L').convert('RGB').save( os.path.join(save_path, 'rand', f'{fname}.png') )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Color reconstructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "save_path = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\ablation\\\\quant_recon'\n",
    "img_size = [256, 256]\n",
    "\n",
    "files = os.listdir(img_path)\n",
    "\n",
    "for f in tqdm(files):\n",
    "    fname = f.split('.')[0]\n",
    "    I_color = Image.open(os.path.join(img_path, f)).convert('RGB').resize(img_size)\n",
    "    I_gray = I_color.convert('L')\n",
    "    idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "    x_color = preprocess(I_color, None).to(args.device)\n",
    "    x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "    # Recon with original color indices\n",
    "    idx_color, idx_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "    idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "    idx_gray = idx_gray.reshape([idx_gray.shape[0]]+idx_shape)\n",
    "    dec_gt = filltran.hybrid_vqgan.decode(idx_color, idx_gray)\n",
    "    dec_gt = output_to_pil(dec_gt[0])\n",
    "\n",
    "    I_color.convert('RGB').save( os.path.join(save_path, 'original', f'{fname}.png') )\n",
    "    dec_gt.convert('RGB').save( os.path.join(save_path, 'recon', f'{fname}.png') )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Final"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\torch\\utils\\tensorboard\\__init__.py:4: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if not hasattr(tensorboard, '__version__') or LooseVersion(tensorboard.__version__) < LooseVersion('1.15'):\n",
      "C:\\Users\\lucky\\Anaconda3\\lib\\site-packages\\torch\\utils\\tensorboard\\__init__.py:4: DeprecationWarning: distutils Version classes are deprecated. Use packaging.version instead.\n",
      "  if not hasattr(tensorboard, '__version__') or LooseVersion(tensorboard.__version__) < LooseVersion('1.15'):\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working with z of shape (1, 1024, 16, 16) = 262144 dimensions.\n"
     ]
    }
   ],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument('--device', type=str, default='cuda:0')\n",
    "\n",
    "parser.add_argument('--root_dir', type=str, default='C:/MyFiles/CondTran/finals/bert_final')\n",
    "parser.add_argument('--log_dir', type=str, default='logs/bert')\n",
    "parser.add_argument('--step', type=str, default='142124')\n",
    "\n",
    "# Load transformer\n",
    "args = parser.parse_args(args=[])\n",
    "os.chdir(args.root_dir)\n",
    "sys.path.append(args.root_dir)\n",
    "module = importlib.import_module(f'filltran.models.colorization')\n",
    "args.fill_model = getattr(module, 'Colorization')\n",
    "filltran = load_model(args.fill_model, args.log_dir, args.step).to(args.device).eval().requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\CondTran\\\\data\\\\raw\\\\in1.JPEG'\n",
    "I_color = Image.open(img_path).convert('RGB')\n",
    "I_gray = I_color.convert('L')\n",
    "idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "x_color = preprocess(I_color, None).to(args.device)\n",
    "x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "# Recon with original color indices\n",
    "idx_color, f_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, f_gray)\n",
    "dec = output_to_pil(dec[0])\n",
    "dec.show()\n",
    "\n",
    "# Recon with full red color indices\n",
    "'''I_allcolor = 128 * np.ones(list(I_color.size[::-1])+[3]).astype(np.uint8)\n",
    "I_allcolor[:, :, 1:3] = 50\n",
    "I_allcolor = Image.fromarray(I_allcolor)\n",
    "x_allcolor = preprocess(I_allcolor, None).to(args.device)\n",
    "idx_allcolor, f_gray = filltran.hybrid_vqgan.encode(x_allcolor, x_gray)\n",
    "idx_allcolor = idx_allcolor.reshape([idx_allcolor.shape[0]]+idx_shape)\n",
    "dec_gray = filltran.hybrid_vqgan.decode(idx_allcolor, f_gray)\n",
    "dec_gray = output_to_pil(dec_gray[0])\n",
    "dec_gray.show()'''\n",
    "\n",
    "# Recon with full gray gray indices\n",
    "'''I_allgray = 128 * np.ones(list(I_color.size[::-1])).astype(np.uint8)\n",
    "I_allgray = Image.fromarray(I_allgray)\n",
    "x_allgray = preprocess(I_allgray, None).to(args.device)\n",
    "idx_color, f_allgray = filltran.hybrid_vqgan.encode(x_color, x_allgray)\n",
    "idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "dec = filltran.hybrid_vqgan.decode(idx_color, f_allgray)\n",
    "dec = output_to_pil(dec[0])\n",
    "dec.show()'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Gray reconstruction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "save_path = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\ablation\\\\chroma_gray'\n",
    "img_size = [256, 256]\n",
    "\n",
    "files = os.listdir(img_path)\n",
    "\n",
    "for f in tqdm(files):\n",
    "    fname = f.split('.')[0]\n",
    "    I_color = Image.open(os.path.join(img_path, f)).convert('RGB').resize(img_size)\n",
    "    I_gray = I_color.convert('L')\n",
    "    idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "    x_color = preprocess(I_color, None).to(args.device)\n",
    "    x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "    # Recon with original color indices\n",
    "    idx_color, f_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "    idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "    dec_gt = filltran.hybrid_vqgan.decode(idx_color, f_gray)\n",
    "    dec_gt = output_to_pil(dec_gt[0])\n",
    "\n",
    "    # Recon with full red color indices\n",
    "    I_rand = np.random.randint(0, 256, size=[1, 1, 3], dtype=np.uint8)\n",
    "    I_rand = I_rand.repeat(I_color.size[1], axis=0).repeat(I_color.size[0], axis=1)\n",
    "    I_rand = Image.fromarray(I_rand)\n",
    "    x_rand = preprocess(I_rand, None).to(args.device)\n",
    "    idx_rand, _ = filltran.hybrid_vqgan.encode(x_rand, None)\n",
    "    idx_rand = idx_rand.reshape([idx_rand.shape[0]]+idx_shape)\n",
    "    dec_rand = filltran.hybrid_vqgan.decode(idx_rand, f_gray)\n",
    "    dec_rand = output_to_pil(dec_rand[0])\n",
    "\n",
    "    I_gray.convert('RGB').save( os.path.join(save_path, 'gray', f'{fname}.png') )\n",
    "    dec_gt.convert('L').convert('RGB').save( os.path.join(save_path, 'gt', f'{fname}.png') )\n",
    "    dec_rand.convert('L').convert('RGB').save( os.path.join(save_path, 'rand', f'{fname}.png') )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Color reconstructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5000/5000 [22:01<00:00,  3.79it/s]\n"
     ]
    }
   ],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "save_path = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\ablation\\\\chroma_recon'\n",
    "img_size = [256, 256]\n",
    "\n",
    "files = os.listdir(img_path)\n",
    "\n",
    "for f in tqdm(files):\n",
    "    fname = f.split('.')[0]\n",
    "    I_color = Image.open(os.path.join(img_path, f)).convert('RGB').resize(img_size)\n",
    "    I_gray = I_color.convert('L')\n",
    "    idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "    x_color = preprocess(I_color, None).to(args.device)\n",
    "    x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "    # Recon with original color indices\n",
    "    idx_color, f_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "    idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "    dec_gt = filltran.hybrid_vqgan.decode(idx_color, f_gray)\n",
    "    dec_gt = output_to_pil(dec_gt[0])\n",
    "\n",
    "    I_color.convert('RGB').save( os.path.join(save_path, 'original', f'{fname}.png') )\n",
    "    dec_gt.convert('RGB').save( os.path.join(save_path, 'recon', f'{fname}.png') )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Gray comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rgb_to_L(x):\n",
    "    x = np.array(x)\n",
    "    lab = cv2.cvtColor(x, cv2.COLOR_RGB2LAB)\n",
    "    l = lab[:, :, 0]\n",
    "\n",
    "    return Image.fromarray(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "img_size = [256, 256]\n",
    "\n",
    "rec_loss = []\n",
    "dec_loss = []\n",
    "lab_loss = []\n",
    "\n",
    "files = os.listdir(img_path)[:10]\n",
    "\n",
    "for f in tqdm(files):\n",
    "    I_color = Image.open(os.path.join(img_path, f)).convert('RGB').resize(img_size)\n",
    "    I_gray = rgb_to_L(I_color)\n",
    "    idx_shape = [I_color.size[1]//16, I_color.size[0]//16]\n",
    "\n",
    "    x_color = preprocess(I_color, None).to(args.device)\n",
    "    x_gray = preprocess(I_gray, None).to(args.device)\n",
    "\n",
    "    # Recon with original color indices\n",
    "    idx_color, f_gray = filltran.hybrid_vqgan.encode(x_color, x_gray)\n",
    "    idx_color = idx_color.reshape([idx_color.shape[0]]+idx_shape)\n",
    "    dec = filltran.hybrid_vqgan.decode(idx_color, f_gray)\n",
    "    dec = output_to_pil(dec[0])\n",
    "    #dec.show()\n",
    "\n",
    "    # Recon with full red color indices\n",
    "    I_allcolor = 128 * np.ones(list(I_color.size[::-1])+[3]).astype(np.uint8)\n",
    "    I_allcolor[:, :, 1:3] = 50\n",
    "    I_allcolor = Image.fromarray(I_allcolor)\n",
    "    x_allcolor = preprocess(I_allcolor, None).to(args.device)\n",
    "    idx_allcolor, _ = filltran.hybrid_vqgan.encode(x_allcolor, None)\n",
    "    idx_allcolor = idx_allcolor.reshape([idx_allcolor.shape[0]]+idx_shape)\n",
    "    dec_gray = filltran.hybrid_vqgan.decode(idx_allcolor, f_gray)\n",
    "    dec_gray = output_to_pil(dec_gray[0])\n",
    "    #dec_gray.show()\n",
    "\n",
    "    # LAB\n",
    "    dec_lab = draw_color(I_color, I_allcolor, [None, None, None, None])\n",
    "    #dec_lab.show()\n",
    "\n",
    "    g = np.array(I_gray)\n",
    "    g_rec = np.array(rgb_to_L(dec_gray))\n",
    "    g_lab = np.array(rgb_to_L(dec_lab))\n",
    "    g_dec = np.array(rgb_to_L(dec))\n",
    "\n",
    "    rec_loss.append( np.sqrt(np.mean(np.square(g_rec - g))) )\n",
    "    dec_loss.append( np.sqrt(np.mean(np.square(g_dec - g))) )\n",
    "    lab_loss.append( np.sqrt(np.mean(np.square(g_lab - g))) )\n",
    "\n",
    "print(np.mean(rec_loss))\n",
    "print(np.mean(dec_loss))\n",
    "print(np.mean(lab_loss))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d2a91ca3d26e2d2b798fe8778458ed9052c39446826a52af3c736ec2371151af"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
